# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: fengyang

import datetime
import logging
import os
import socket

from absl import app, flags
FLAGS = flags.FLAGS


def raname_tpm_files_in_dir(feed_dir):
  files = os.listdir(feed_dir)
  now = datetime.datetime.now()
  for filename in files:
    if not filename.startswith('_'):
      continue
    full_path = os.path.join(feed_dir, filename)
    m_time = datetime.datetime.fromtimestamp(os.path.getmtime(full_path))
    interval = now - m_time
    # write a file every 30 minutes
    safe_interval = datetime.timedelta(minutes=40)
    if interval > safe_interval:
      tmp_name = full_path
      new_name = os.path.join(feed_dir, filename[1:])
      logging.info("Rename %s to %s" % (full_path, new_name))
      os.rename(tmp_name, new_name)


def rename_tmp_files(flow_root_dir=None, machine=None, date=None):
  if flow_root_dir is None:
    flow_root_dir = "/data/feed-data"
  if machine is None:
    machine = socket.gethostname()

  if date is None:
    # no specified, check 2 days
    today = datetime.date.today()
    yesterday = datetime.date.today() - datetime.timedelta(days=1)
    today_str = today.strftime('%Y%m%d')
    yesterday_str = yesterday.strftime('%Y%m%d')
    feed_dirs = [
        os.path.join(flow_root_dir, machine, today_str),
        os.path.join(flow_root_dir, machine, yesterday_str)
    ]
  else:
    feed_dirs = [os.path.join(flow_root_dir, machine, date)]

  for feed_dir in feed_dirs:
    raname_tpm_files_in_dir(feed_dir)
    sub_dirs = next(os.walk(feed_dir))[1]
    sub_dirs = [os.path.join(feed_dir, sub_dir) for sub_dir in sub_dirs]
    for sub_dir in sub_dirs:
      raname_tpm_files_in_dir(sub_dir)


def main(_):
  rename_tmp_files(FLAGS.flow_root_dir, FLAGS.machine, FLAGS.date)


if __name__ == '__main__':
  flags.DEFINE_string('flow_root_dir', None, 'flow root dir')
  flags.DEFINE_string('machine', None, 'hostname')
  flags.DEFINE_string('date', None, 'yyyymmdd')

  app.run(main)
